# Copyright 2016-2024, Pulumi Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import traceback
import uuid
from collections.abc import Awaitable, Callable, Mapping
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
    Union,
    cast,
)

import grpc
from google.protobuf.message import Message
from grpc import aio
from grpc.experimental.aio import init_grpc_aio

from .. import log
from ..invoke import InvokeOptions, InvokeTransform
from .proto import (
    alias_pb2,
    callback_pb2,
    callback_pb2_grpc,
    resource_pb2,
    resource_pb2_grpc,
)
from .rpc import (
    deserialize_properties,
    is_rpc_secret,
    serialize_properties,
    unwrap_rpc_secret,
)
from ._grpc_settings import _GRPC_CHANNEL_OPTIONS

if TYPE_CHECKING:
    from ..resource import (
        Alias,
        ResourceOptions,
        ResourceTransform,
    )
    from ..resource_hooks import ResourceHook


# Workaround for https://github.com/grpc/grpc/issues/38679,
# https://github.com/grpc/grpc/issues/22365#issuecomment-2254278769
# This will be fixed in grpcio 1.75.1 with https://github.com/grpc/grpc/pull/40649
init_grpc_aio()

_CallbackFunction = Callable[[bytes], Awaitable[Message]]


class _CallbackServicer(callback_pb2_grpc.CallbacksServicer):
    _servicers: list[_CallbackServicer] = []

    _callbacks: dict[str, _CallbackFunction]
    _monitor: resource_pb2_grpc.ResourceMonitorStub
    _server: aio.Server
    _target: str

    _transforms: dict[Union[ResourceTransform, InvokeTransform], str]

    def __init__(self, monitor: resource_pb2_grpc.ResourceMonitorStub):
        log.debug("Creating CallbackServicer")
        _CallbackServicer._servicers.append(self)
        self._callbacks = {}
        self._transforms = {}
        self._monitor = monitor
        self._server = aio.server(options=_GRPC_CHANNEL_OPTIONS)
        callback_pb2_grpc.add_CallbacksServicer_to_server(self, self._server)
        port = self._server.add_insecure_port(address="127.0.0.1:0")
        self._target = f"127.0.0.1:{port}"

    async def serve(self):
        await self._server.start()

    @classmethod
    async def shutdown(cls):
        for servicer in cls._servicers:
            await servicer._server.stop(grace=0)

    async def Invoke(
        self,
        request: callback_pb2.CallbackInvokeRequest,
        context: aio.ServicerContext[
            callback_pb2.CallbackInvokeRequest, callback_pb2.CallbackInvokeResponse
        ],
    ) -> callback_pb2.CallbackInvokeResponse:
        log.debug(f"Invoke callback {request.token}")
        callback = self._callbacks.get(request.token)
        if callback is None:
            await context.abort(
                grpc.StatusCode.INVALID_ARGUMENT,
                f"Callback with token {request.token} not found!",
            )
            raise Exception("Unreachable, abort never returns")

        try:
            response = await callback(request.request)
            return callback_pb2.CallbackInvokeResponse(
                response=response.SerializeToString()
            )
        except Exception as e:  # noqa: BLE001 catch blind exception
            log.debug(
                f"Invoke callback {request.token} error: {traceback.format_exc()}"
            )
            details = self.pretty_error(e)
            await context.abort(
                grpc.StatusCode.UNKNOWN,
                f"Invoke callback {request.token} error: {details}",
            )
        except asyncio.CancelledError as e:
            log.debug(
                f"Invoke callback {request.token} was cancelled: {traceback.format_exc()}"
            )
            details = self.pretty_error(e)
            await context.abort(
                grpc.StatusCode.CANCELLED,
                f"Callback with token {request.token} was cancelled: {details}",
            )

    def pretty_error(self, e: BaseException) -> str:
        """
        pretty_error filters a stack trace down to user code and formats it for display
        """
        stack = traceback.extract_tb(e.__traceback__)
        # Drop internal stack frames
        filtered_stack = [f for f in stack if f.filename != __file__]
        # If there are no stack frames left after dropping internal frames,
        # use the original stack, this is likely an internal error.
        filtered_stack = filtered_stack if len(filtered_stack) > 0 else stack
        pretty_stack = "".join(traceback.format_list(filtered_stack))
        details = f"{str(e)}:\n{pretty_stack}"
        return details

    def register_transform(self, transform: ResourceTransform) -> callback_pb2.Callback:
        # If this transform function has already been registered, return it.
        token = self._transforms.get(transform)
        if token is not None:
            return callback_pb2.Callback(token=token, target=self._target)

        from ..resource import (
            ResourceTransformArgs,
            ResourceTransformResult,
        )

        async def cb(s: bytes) -> Message:
            request: resource_pb2.TransformRequest = (
                resource_pb2.TransformRequest.FromString(s)
            )

            args = ResourceTransformArgs(
                custom=request.custom,
                type_=request.type,
                name=request.name,
                props=deserialize_properties(request.properties),
                opts=self._resource_options(request),
            )

            maybeAwaitable = transform(args)
            result: Optional[ResourceTransformResult] = None
            if isinstance(maybeAwaitable, Awaitable):
                result = await maybeAwaitable
            else:
                result = maybeAwaitable

            if result is None:
                return resource_pb2.TransformResponse(
                    properties=request.properties,
                    options=request.options,
                )

            result_props = await serialize_properties(result.props, {})

            result_opts = (
                await self._transformation_resource_options(result.opts, request.name)
                if result.opts is not None
                else None
            )

            return resource_pb2.TransformResponse(
                properties=result_props,
                options=result_opts,
            )

        token = str(uuid.uuid4())
        self._callbacks[token] = cb
        self._transforms[transform] = token
        return callback_pb2.Callback(
            token=token,
            target=self._target,
        )

    def register_stack_transform(self, transform: ResourceTransform):
        callback = self.register_transform(transform)
        try:
            self._monitor.RegisterStackTransform(callback)
        except:
            # Remove the transform since we didn't manage to actually register it.
            self._transforms.pop(transform)
            self._callbacks.pop(callback.token)
            raise

    def do_register_invoke_transform(
        self, transform: InvokeTransform
    ) -> callback_pb2.Callback:
        # If this transform function has already been registered, return it.
        token = self._transforms.get(transform)
        if token is not None:
            return callback_pb2.Callback(token=token, target=self._target)

        from ..invoke import (
            InvokeTransformArgs,
            InvokeTransformResult,
        )

        async def cb(s: bytes) -> Message:
            request: resource_pb2.TransformInvokeRequest = (
                resource_pb2.TransformInvokeRequest.FromString(s)
            )

            args = InvokeTransformArgs(
                token=request.token,
                args=deserialize_properties(request.args),
                opts=self._invoke_options(request),
            )

            maybeAwaitable = transform(args)
            result: Optional[InvokeTransformResult] = None
            if isinstance(maybeAwaitable, Awaitable):
                result = await maybeAwaitable
            else:
                result = maybeAwaitable

            if result is None:
                return resource_pb2.TransformInvokeResponse(
                    args=request.args,
                    options=request.options,
                )

            result_args = await serialize_properties(result.args, {})

            result_opts = (
                await self._transformation_invoke_options(result.opts)
                if result.opts is not None
                else None
            )

            return resource_pb2.TransformInvokeResponse(
                args=result_args,
                options=result_opts,
            )

        token = str(uuid.uuid4())
        self._callbacks[token] = cb
        self._transforms[transform] = token
        return callback_pb2.Callback(
            token=token,
            target=self._target,
        )

    def register_invoke_transform(self, transform: InvokeTransform):
        callback = self.do_register_invoke_transform(transform)
        try:
            self._monitor.RegisterStackInvokeTransform(callback)
        except:
            # Remove the transform since we didn't manage to actually register it.
            self._transforms.pop(transform)
            self._callbacks.pop(callback.token)
            raise

    def do_register_resource_hook(
        self, hook: ResourceHook
    ) -> resource_pb2.RegisterResourceHookRequest:
        log.debug(f"do_register_resource_hook {hook.name}")

        async def cb(s: bytes) -> Message:
            request: resource_pb2.ResourceHookRequest = (
                resource_pb2.ResourceHookRequest.FromString(s)
            )

            try:
                from ..resource_hooks import ResourceHookArgs

                new_inputs = (
                    outputtify_secrets(
                        cast(
                            dict[str, Any],
                            deserialize_properties(
                                request.new_inputs,
                                keep_unknowns=True,
                            ),
                        )
                    )
                    if request.new_inputs
                    else None
                )
                old_inputs = (
                    outputtify_secrets(
                        cast(
                            dict[str, Any],
                            deserialize_properties(
                                request.old_inputs,
                                keep_unknowns=True,
                            ),
                        )
                    )
                    if request.old_inputs
                    else None
                )
                new_outputs = (
                    outputtify_secrets(
                        cast(
                            dict[str, Any],
                            deserialize_properties(
                                request.new_outputs,
                                keep_unknowns=True,
                            ),
                        )
                    )
                    if request.new_outputs
                    else None
                )
                old_outputs = (
                    outputtify_secrets(
                        cast(
                            dict[str, Any],
                            deserialize_properties(
                                request.old_outputs,
                                keep_unknowns=True,
                            ),
                        )
                    )
                    if request.old_outputs
                    else None
                )

                maybeAwaitable = hook(
                    ResourceHookArgs(
                        urn=request.urn,
                        id=request.id,
                        name=request.name,
                        type=request.type,
                        new_inputs=new_inputs,
                        old_inputs=old_inputs,
                        new_outputs=new_outputs,
                        old_outputs=old_outputs,
                    )
                )
                if isinstance(maybeAwaitable, Awaitable):
                    await maybeAwaitable
                return resource_pb2.ResourceHookResponse()
            except Exception as e:  # noqa: BLE001 catch blind exception
                log.debug(f"Exception while executing hook: {str(e)}")
                return resource_pb2.ResourceHookResponse(error=str(e))

        token = str(uuid.uuid4())
        self._callbacks[token] = cb
        callback = callback_pb2.Callback(
            token=token,
            target=self._target,
        )
        return resource_pb2.RegisterResourceHookRequest(
            name=hook.name,
            callback=callback,
            on_dry_run=hook.opts.on_dry_run if hook.opts else False,
        )

    def register_resource_hook(self, hook: ResourceHook) -> None:
        req = self.do_register_resource_hook(hook)
        try:
            self._monitor.RegisterResourceHook(req)
        except:
            # Remove the hook since we didn't manage to actually register it.
            self._callbacks.pop(req.callback.token)
            raise

    def _resource_options(
        self, request: resource_pb2.TransformRequest
    ) -> ResourceOptions:
        from ..resource import (
            CustomTimeouts,
            DependencyProviderResource,
            DependencyResource,
            ResourceOptions,
        )

        opts = (
            request.options
            if request.HasField("options")
            else resource_pb2.TransformResourceOptions()
        )

        ropts = ResourceOptions()

        if opts.HasField("delete_before_replace"):
            ropts.delete_before_replace = opts.delete_before_replace

        id = getattr(opts, "import", None)
        if id:
            ropts.import_ = id

        additional_secret_outputs = list(opts.additional_secret_outputs)
        if additional_secret_outputs:
            ropts.additional_secret_outputs = additional_secret_outputs

        providers = {
            k: DependencyProviderResource(v) for k, v in opts.providers.items()
        }
        if providers:
            ropts.providers = providers

        aliases = [self._alias(a) for a in opts.aliases]
        if aliases:
            ropts.aliases = aliases

        if opts.HasField("custom_timeouts"):
            custom_timeouts = CustomTimeouts()
            if opts.custom_timeouts.create:
                custom_timeouts.create = opts.custom_timeouts.create
            if opts.custom_timeouts.update:
                custom_timeouts.update = opts.custom_timeouts.update
            if opts.custom_timeouts.delete:
                custom_timeouts.delete = opts.custom_timeouts.delete
            ropts.custom_timeouts = custom_timeouts

        if opts.deleted_with:
            ropts.deleted_with = DependencyResource(opts.deleted_with)

        depends_on = [DependencyResource(d) for d in opts.depends_on]
        if depends_on:
            ropts.depends_on = depends_on

        ignore_changes = list(opts.ignore_changes)
        if ignore_changes:
            ropts.ignore_changes = ignore_changes

        if request.parent:
            ropts.parent = DependencyResource(request.parent)

        if opts.plugin_download_url:
            ropts.plugin_download_url = opts.plugin_download_url

        if opts.HasField("protect"):
            ropts.protect = opts.protect

        if opts.provider:
            ropts.provider = DependencyProviderResource(opts.provider)

        replace_on_changes = list(opts.replace_on_changes)
        if replace_on_changes:
            ropts.replace_on_changes = replace_on_changes

        if opts.retain_on_delete:
            ropts.retain_on_delete = opts.retain_on_delete

        if opts.version:
            ropts.version = opts.version

        if opts.hooks:
            from .resource_hooks import _binding_from_proto

            ropts.hooks = _binding_from_proto(opts.hooks)

        return ropts

    def _alias(self, alias: alias_pb2.Alias) -> Union[str, Alias]:
        if alias.HasField("urn"):
            return alias.urn

        from ..resource import Alias

        a = Alias()
        if alias.spec.name:
            a.name = alias.spec.name
        if alias.spec.type:
            a.type_ = alias.spec.type
        if alias.spec.project:
            a.project = alias.spec.project
        if alias.spec.stack:
            a.stack = alias.spec.stack
        if alias.spec.parentUrn:
            a.parent = alias.spec.parentUrn
        elif alias.spec.noParent:
            a.parent = None
        return a

    async def _transformation_resource_options(
        self,
        opts: ResourceOptions,
        resource_name: str,
    ) -> resource_pb2.TransformResourceOptions:
        from ._depends_on import (
            _resolve_depends_on_urns,
        )
        from .resource import (
            _create_custom_timeouts,
            _create_provider_ref,
            create_alias_spec,
            _prepare_resource_hooks,
        )
        from ..resource import (
            Alias,
            ProviderResource,
            _collapse_providers,
        )
        from ..output import Output

        aliases: list[alias_pb2.Alias] = []
        if opts.aliases:
            for alias in opts.aliases:
                resolved = await Output.from_input(alias).future()
                if resolved is None:
                    continue
                if isinstance(resolved, str):
                    aliases.append(alias_pb2.Alias(urn=resolved))
                else:
                    spec = await create_alias_spec(cast(Alias, resolved))
                    aliases.append(alias_pb2.Alias(spec=spec))

        custom_timeouts = None
        if opts.custom_timeouts is not None:
            custom_timeouts = _create_custom_timeouts(opts.custom_timeouts)

        depends_on = await _resolve_depends_on_urns(opts._depends_on_list())

        ignore_changes = None
        if opts.ignore_changes:
            ignore_changes = list(opts.ignore_changes)

        replace_on_changes = None
        if opts.replace_on_changes:
            replace_on_changes = list(opts.replace_on_changes)

        additional_secret_outputs = None
        if opts.additional_secret_outputs:
            additional_secret_outputs = list(opts.additional_secret_outputs)

        hooks = await _prepare_resource_hooks(opts.hooks, resource_name)

        result = resource_pb2.TransformResourceOptions(
            aliases=aliases or None,
            custom_timeouts=custom_timeouts,
            depends_on=depends_on or None,
            ignore_changes=ignore_changes,
            replace_on_changes=replace_on_changes,
            additional_secret_outputs=additional_secret_outputs,
            hooks=hooks,
        )

        if opts.import_ is not None:
            setattr(result, "import", opts.import_)
        if opts.deleted_with is not None:
            result.deleted_with = cast(str, await opts.deleted_with.urn.future())
        if opts.replace_with:
            for dependency in opts.replace_with:
                urn = await dependency.urn.future()
                if urn:
                    result.replace_with.append(cast(str, urn))
        if opts.plugin_download_url:
            result.plugin_download_url = opts.plugin_download_url
        if opts.protect is not None:
            result.protect = opts.protect
        if opts.retain_on_delete:
            result.retain_on_delete = opts.retain_on_delete
        if opts.version:
            result.version = opts.version
        if opts.provider is not None:
            result.provider = await _create_provider_ref(opts.provider)
        if opts.delete_before_replace:
            result.delete_before_replace = opts.delete_before_replace
        if opts.providers:
            _collapse_providers(opts)
            providers = cast(Mapping[str, ProviderResource], opts.providers)
            result.providers.update(
                {k: await _create_provider_ref(v) for k, v in providers.items()}
            )

        return result

    def _invoke_options(
        self, request: resource_pb2.TransformInvokeRequest
    ) -> InvokeOptions:
        from ..resource import (
            DependencyProviderResource,
        )

        opts = (
            request.options
            if request.HasField("options")
            else resource_pb2.TransformInvokeOptions()
        )

        ropts = InvokeOptions()

        if opts.plugin_download_url:
            ropts.plugin_download_url = opts.plugin_download_url

        if opts.provider:
            ropts.provider = DependencyProviderResource(opts.provider)

        if opts.version:
            ropts.version = opts.version

        return ropts

    async def _transformation_invoke_options(
        self, opts: InvokeOptions
    ) -> resource_pb2.TransformInvokeOptions:
        from .resource import (
            _create_provider_ref,
        )

        result = resource_pb2.TransformInvokeOptions()

        if opts.plugin_download_url:
            result.plugin_download_url = opts.plugin_download_url
        if opts.version:
            result.version = opts.version
        if opts.provider is not None:
            result.provider = await _create_provider_ref(opts.provider)

        return result


def outputtify_secrets(props: dict[str, Any]) -> dict[str, Any]:
    from ..output import Output

    for k, v in props.items():
        if is_rpc_secret(v):
            props[k] = Output.secret(unwrap_rpc_secret(v))

    return props
